#!/usr/bin/env python3
"""
数据库迁移脚本：添加AI提示词配置表
"""

import sqlite3
from pathlib import Path
import sys

# 默认的Email提示词模板
DEFAULT_EMAIL_PROMPT = """你是一位专业的邮件内容分析师。你的任务是精准地总结任何类型的邮件，无论是个人通讯、工作往来，还是资讯订阅（如新闻、科技、产品更新等）。
请遵循以下指南，并使用**美观、易读的Markdown格式**输出总结：
1.  **识别邮件核心目的**：快速判断邮件的主要意图（是通知、请求、确认，还是资讯分享？）。
2.  **提取关键信息**：
    *   **对于普通邮件**：明确总结出发件人的核心诉求、需要你执行的行动（To-Dos）、以及重要的时间节点或截止日期。
    *   **对于订阅/资讯类邮件**：提炼出主要观点、核心事件、关键数据或新发布的功能。使用要点列表（bullet points）来呈现，使其清晰明了。
3.  **格式化输出**：
    *   使用 `#` 或 `##` 作为主标题，清晰概括邮件主题。
    *   使用加粗（`**`）突出显示人名、项目名、日期和关键术语。
    *   使用项目符号（`-` 或 `*`）来组织信息点，保持结构清晰。
    *   在不同信息模块之间适当使用空行，增加可读性。
    *   总结内容应保持客观、中立，忠于原文，不要添加个人评论或推断。
请对以下邮件内容进行总结：

邮件主题：{subject}
发件人：{sender}
邮件内容：
{content}

请用中文回复，格式清晰，重点突出。"""


def migrate_database():
    """执行迁移：添加ai_prompt_configs表"""
    db_path = Path(__file__).parent.parent / "instance" / "feng_admin.db"
    
    if not db_path.exists():
        print(f"❌ 数据库文件不存在: {db_path}")
        sys.exit(1)
    
    print(f"📁 数据库路径: {db_path}")
    
    conn = sqlite3.connect(str(db_path))
    cursor = conn.cursor()
    
    try:
        # 检查表是否已存在
        cursor.execute("""
            SELECT name FROM sqlite_master 
            WHERE type='table' AND name='ai_prompt_configs'
        """)
        
        if cursor.fetchone():
            print("⚠️  表 ai_prompt_configs 已存在，跳过创建")
        else:
            # 创建ai_prompt_configs表
            cursor.execute("""
                CREATE TABLE ai_prompt_configs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id INTEGER NOT NULL,
                    module VARCHAR(50) NOT NULL,
                    prompt_template TEXT NOT NULL,
                    is_active BOOLEAN DEFAULT 1,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY (user_id) REFERENCES users(id),
                    UNIQUE (user_id, module)
                )
            """)
            print("✅ 创建表 ai_prompt_configs 成功")
        
        # 为所有现有用户插入默认的email提示词配置
        cursor.execute("SELECT id FROM users")
        users = cursor.fetchall()
        
        inserted_count = 0
        skipped_count = 0
        
        for (user_id,) in users:
            # 检查该用户是否已有email配置
            cursor.execute("""
                SELECT id FROM ai_prompt_configs 
                WHERE user_id = ? AND module = 'email'
            """, (user_id,))
            
            if cursor.fetchone():
                skipped_count += 1
                continue
            
            # 插入默认配置
            cursor.execute("""
                INSERT INTO ai_prompt_configs (user_id, module, prompt_template, is_active)
                VALUES (?, 'email', ?, 1)
            """, (user_id, DEFAULT_EMAIL_PROMPT))
            inserted_count += 1
        
        conn.commit()
        
        print(f"✅ 迁移成功完成！")
        print(f"   - 为 {inserted_count} 个用户创建了默认email提示词配置")
        if skipped_count > 0:
            print(f"   - 跳过 {skipped_count} 个已存在的配置")
        
    except Exception as e:
        conn.rollback()
        print(f"❌ 迁移失败: {e}")
        sys.exit(1)
    finally:
        conn.close()


def rollback_migration():
    """回滚迁移：删除ai_prompt_configs表"""
    db_path = Path(__file__).parent.parent / "instance" / "feng_admin.db"
    
    if not db_path.exists():
        print(f"❌ 数据库文件不存在: {db_path}")
        sys.exit(1)
    
    print(f"📁 数据库路径: {db_path}")
    print("⚠️  警告：即将删除 ai_prompt_configs 表及其所有数据！")
    
    confirm = input("确认回滚？(yes/no): ")
    if confirm.lower() != 'yes':
        print("❌ 回滚已取消")
        return
    
    conn = sqlite3.connect(str(db_path))
    cursor = conn.cursor()
    
    try:
        cursor.execute("DROP TABLE IF EXISTS ai_prompt_configs")
        conn.commit()
        print("✅ 回滚成功：已删除 ai_prompt_configs 表")
    except Exception as e:
        conn.rollback()
        print(f"❌ 回滚失败: {e}")
        sys.exit(1)
    finally:
        conn.close()


def check_migration_status():
    """检查迁移状态"""
    db_path = Path(__file__).parent.parent / "instance" / "feng_admin.db"
    
    if not db_path.exists():
        print(f"❌ 数据库文件不存在: {db_path}")
        sys.exit(1)
    
    conn = sqlite3.connect(str(db_path))
    cursor = conn.cursor()
    
    try:
        # 检查表是否存在
        cursor.execute("""
            SELECT name FROM sqlite_master 
            WHERE type='table' AND name='ai_prompt_configs'
        """)
        
        if cursor.fetchone():
            print("✅ ai_prompt_configs 表已存在")
            
            # 统计配置数量
            cursor.execute("""
                SELECT module, COUNT(*) as count
                FROM ai_prompt_configs
                GROUP BY module
            """)
            
            configs = cursor.fetchall()
            if configs:
                print("\n📊 提示词配置统计：")
                for module, count in configs:
                    print(f"   - {module}: {count} 个配置")
            else:
                print("   (暂无配置)")
        else:
            print("❌ ai_prompt_configs 表不存在，需要执行迁移")
    
    except Exception as e:
        print(f"❌ 检查失败: {e}")
        sys.exit(1)
    finally:
        conn.close()


if __name__ == '__main__':
    import sys
    
    if len(sys.argv) > 1:
        command = sys.argv[1]
        if command == 'rollback':
            rollback_migration()
        elif command == 'check':
            check_migration_status()
        else:
            print("用法:")
            print("  python migrate_add_ai_prompt_config.py        # 执行迁移")
            print("  python migrate_add_ai_prompt_config.py check  # 检查状态")
            print("  python migrate_add_ai_prompt_config.py rollback  # 回滚")
    else:
        migrate_database()

